To my mind, we would get better error handling (eg stack traces) if we resolved the __torch_dispatch__ earlier by checking for subclasses in the torchsymbol handling logic. Also, I don't think we have much information about the outputs - e.g. are they subclasses again or the original class, shapes etc. - without doing so, so we cannot evaluate even metadata in control flow.

There's a lot going on with this PR, and it's pretty complicated. Maybe we should schedule an online sync, @crcrpar and @IvanYashchuk, to see if we can make it more incremental?

I would still prefer to do the flattening during interpretation for the benefit of getting error messages with backtraces.

There's a lot going on with this PR, and it's pretty complicated. Maybe we should schedule an online sync, @crcrpar and @IvanYashchuk, to see if we can make it more incremental?

We don't have a nice time slots that work all of us.
What would have you find PRs incremental?
Would it look better to you to split this into two, one to add a proxy and the other to add trace transforms?

EDIT: #1583, #1584, and #1585 are a sequence of PRs that cover this one and #1415.

I would still prefer to do the flattening during interpretation for the benefit of getting error messages with backtraces.

I embarrassingly am not familiar with interpretation implementation at all. Could you give me some pointers to add flattening so that it happens during interpretation?

EDIT: __torch_dispatch__ overrides the behavior of ops used in backward and the customized behavior could use torch ops that do not have their backward definition. Thus I do think we'd need to put the flattening AFTER forward-backward split as in #1415

There's a lot going on with this PR, and it's pretty complicated. Maybe we should schedule an online sync, @crcrpar and @IvanYashchuk, to see if we can make it more incremental?

We don't have a nice time slots that work all of us. What would have you find PRs incremental? Would it look better to you to split this into two, one to add a proxy and the other to add trace transforms?

OK; we can try to work asynchronously. For a first incremental PR, would you create a PR adding support for aten operators? In particular, if someone were to call something like torch.ops.aten.add then how would that work? Can that operator be added to the torch operations (and be in the torch language context), or should it be added to its own aten language context that has a separate file (or set of files)?

if someone were to call something like torch.ops.aten.add then how would that work?

So far in my implementation there's an optimistic mapping from core aten ops to ltorch ops:

By the way, if we're to cover core aten ops, then I'd say it'd be worth thinking of using thunder as a custom backend after AOTAutograd.

Can that operator be added to the torch operations (and be in the torch language context), or should it be added to its own aten language context that has a separate file (or set of files)?

Currently torchsymbol has some for core aten ops, apparently

if someone were to call something like torch.ops.aten.add then how would that work?

So far in my implementation there's an optimistic mapping from core aten ops to ltorch ops:

lightning-thunder/thunder/transforms/tensor_wrapper_subclass.py

Lines 339 to 348 in 515d425

	for node in list_of_function_call_node:
	if not hasattr(ltorch, node.target._opname):
	msg = (
	f"`thunder.torch` does not have corresponding op for {node.target._opname}. "
	"Think about adding it to thunder/torch/default_torch_ops.py"
	f"\nThe op is found while flattening the following BoundSymbol:\n{bsym}"
	f"\ntorch.fx graph:\n{fx_graph.print_readable(print_output=False)}"
	)
	raise RuntimeError(msg)
	ltorch_ops_for_node_of_ops.append(getattr(ltorch, node.target._opname))

.
By the way, if we're to cover core aten ops, then I'd say it'd be worth thinking of using thunder as a custom backend after AOTAutograd.

Can that operator be added to the torch operations (and be in the torch language context), or should it be added to its own aten language context that has a separate file (or set of files)?

Currently torchsymbol has some for core aten ops, apparently

lightning-thunder/thunder/torch/__init__.py

Lines 172 to 173 in 9d79b8d

	elif hasattr(torch.ops.aten, name):
	id = f"torch.ops.aten.{name}"

lightning-thunder/thunder/torch/__init__.py

Line 4700 in 9d79b8d

@torchsymbol(torch.ops.aten.embedding_backward)

and

lightning-thunder/thunder/torch/__init__.py

Line 4308 in 9d79b8d

torch.ops.aten._adaptive_avg_pool2d_backward,

are a core aten op. So I think extending thunder/torch/__init__.py would be fair.

OK; expanding thunder/torch/init.py sounds good for now. Let's not "optimistically" try to map ATen operations to torch operations for the moment, but just treat them like different operations.

Would you submit a PR adding torch.ops.aten.add to the torch operations?

EDITED BELOW.

As a follow-up PR to that, what about working with a program like

# Original program
def foo(x):
  return x

# Trace
def computation(x):
  # x: "MyTensorSubclass[cuda:0 f32[12, 12]]" 
  return x

Where the initial trace shows the tensor subclass and its flattened information, and the prologue validates the subclass and its flattening. Then I'd be curious to see addition with that tensor, like this:

# Original program
def foo(x):
  return x + 1

# Trace
def computation(x):
  # x: "MyTensorSubclass[cuda:0 f32[12, 12]]" 
  t0 = MyTensorSubclass.torch.add(x, 1)  # t0: "MyTensorSubclass[cuda:0 f32[12, 12]]"
    # t1 = flatten_tensor_subclass(MyTensorSubclass, x)
    # t2 = torch.ops.aten.add(t1, 1)
      # <decomposition of aten.add into prims would go here> 
    # t0 = unflatten_tensor_subclass(MyTensorSubclass, t2)
  return t0

This can be translated for execution by PyTorch, but I think working through this will be interesting. Then the follow-up question is what the grad transform for it looks like, and how this operation should be translated for execution by nvFuser.

IMHO, it'd sound more natural to me to register core aten ops to thunder.torch namespace after merging #1583.

Then registration comes after the aforementioned PR, before #1584 and #1585, followed by some refinement of prologue and how traces with tensor subclasses look accompanied by #1584.

the follow-up question is what the grad transform for it looks like, and how this operation should be translated for execution by nvFuser.

With the experience of #1585, I do think we'd have to let the trace get split into forward and backward before interpreting __torch_dispatch__ partly because the extended behavior of certain ops could be dependent on ops without any backward definitions.

IMHO, it'd sound more natural to me to register core aten ops to thunder.torch namespace after merging #1583.

Maybe? I guess it just seems like an easy first step to understand the application of aten operators.

Then registration comes after the aforementioned PR, before #1584 and #1585, followed by some refinement of prologue and how traces with tensor subclasses look accompanied by #1584.

the follow-up question is what the grad transform for it looks like, and how this operation should be translated for execution by nvFuser.

With the experience of #1585, I do think we'd have to let the trace get split into forward and backward before interpreting __torch_dispatch__ partly because the extended behavior of certain ops could be dependent on ops without any backward definitions.

Could you elaborate on this? It would be helpful to see an example. I was thinking that a program like

# Trace
def computation(x):
  # x: "MyTensorSubclass[cuda:0 f32[12, 12]]" 
  t0 = MyTensorSubclass.torch.add(x, 1)  # t0: "MyTensorSubclass[cuda:0 f32[12, 12]]"
    # t1 = flatten_tensor_subclass(MyTensorSubclass, x)
    # t2 = torch.ops.aten.add(t1, 1)
      # <decomposition of aten.add into prims would go here> 
    # t0 = unflatten_tensor_subclass(MyTensorSubclass, t2)
  return t0

Could be auto-differentiated like programs today because MyTensorSubclass.torch.add would not define a grad transform, and so the grad transform would "flatten" it to its components, which would be:

# t1 = flatten_tensor_subclass(MyTensorSubclass, x)
    # t2 = torch.ops.aten.add(t1, 1)
      # <decomposition of aten.add into prims would go here> 
    # t0 = unflatten_tensor_subclass(MyTensorSubclass, t2)

Which I expect would get further flattened into a prims.add call. Maybe @IvanYashchuk has a different perspective.

Because of that I was thinking that we'd want to understand the tensor subclass operations ASAP. What are your thoughts, @crcrpar?

I choose to apply the trace transform AFTER the forward-backward split after I worked on torchao.float8 linears.

Full traces are available on https://gist.github.com/crcrpar/9b69ef83e68e306415af091d025cbf9c.
A program is (quote of https://gist.github.com/crcrpar/9b69ef83e68e306415af091d025cbf9c#file-0_torchao_fp8linear-py)

import torch.nn as nn
from torchao.float8 import convert_to_float8_training

def main():
    batch_size, in_features, out_features = 16, 32, 64
    device, dtype = torch.device("cuda"), torch.float32
    fp8_model = convert_to_float8_training(nn.Linear(in_features, out_features).to(device=device, dtype=dtype))
    jitted = thunder.jit(fp8_model)

The first computation trace has the following lines (= thunder.last_traces(jitted)[0], full trace is https://gist.github.com/crcrpar/9b69ef83e68e306415af091d025cbf9c#file-1_0_first_fwd_trace-py)

  # /path/to/site-packages/torchao/float8/float8_linear.py:57: 	        res_bits = res_bits.reshape(*orig_shape[:-1], res_bits.shape[-1])
  tensor = manual_float8_matmul_with_args_in_float8_140062104971168_2(input_fp8, weight_fp8_t)  # tensor: "cuda:0 f32[16, 64]"
    # t91 = ltorch.reshape(input_fp8, -1, 32)  # t91: "cuda:0 f32[16, 32]"
      # t91 = prims.reshape(input_fp8, (16, 32))  # t91: "cuda:0 f32[16, 32]"
    # tensor = ltorch.spmm(t91, weight_fp8_t)  # tensor: "cuda:0 f32[16, 64]"

The corresponding lines are https://github.com/pytorch/ao/blob/fe5f11b2c58b452e01ba9ec7359629928b143619/torchao/float8/float8_linear.py#L32-L58, which is the forward of manual_float8_matmul_with_args_in_float8(torch.autograd.Function).

Here, t91 and weight_fp8_t are an object of torchao.float8.Float8Tensor.

Float8Tensor.__torch_dispatch__ defines its custom behavior for mm which calls torch._scaled_mm as in https://github.com/pytorch/ao/blob/fe5f11b2c58b452e01ba9ec7359629928b143619/torchao/float8/float8_python_api.py#L22.

These lines are rewritten to

  # [3/8] unrolled `__torch_dispatch__` of `torch.spmm(t212, t205)`
  t214 = ltorch.core_aten_t(t204)  # t214: "cuda:0 f8_e4m3fn[64, 32]"
    # t214 = ltorch.transpose(t204, 0, 1)  # t214: "cuda:0 f8_e4m3fn[64, 32]"
      # t214 = prims.transpose(t204, (1, 0))  # t214: "cuda:0 f8_e4m3fn[64, 32]"
  # [4/8] unrolled `__torch_dispatch__` of `torch.spmm(t212, t205)`
  t215 = ltorch.core_aten_clone(t214, memory_format=None)  # t215: "cuda:0 f8_e4m3fn[64, 32]"
    # t215 = prims.clone(t214)  # t215: "cuda:0 f8_e4m3fn[64, 32]"
  # [5/8] unrolled `__torch_dispatch__` of `torch.spmm(t212, t205)`
  t216 = ltorch.core_aten_t(t215)  # t216: "cuda:0 f8_e4m3fn[32, 64]"
    # t216 = ltorch.transpose(t215, 0, 1)  # t216: "cuda:0 f8_e4m3fn[32, 64]"
      # t216 = prims.transpose(t215, (1, 0))  # t216: "cuda:0 f8_e4m3fn[32, 64]"
  # [6/8] unrolled `__torch_dispatch__` of `torch.spmm(t212, t205)`
  t217 = ltorch.core_aten_reciprocal(scale)  # t217: "cuda:0 f32[]"
    # t217 = prims.reciprocal(scale)  # t217: "cuda:0 f32[]"
  # [7/8] unrolled `__torch_dispatch__` of `torch.spmm(t212, t205)`
  t218 = ltorch.core_aten_reciprocal(weight_scale)  # t218: "cuda:0 f32[]"
    # t218 = prims.reciprocal(weight_scale)  # t218: "cuda:0 f32[]"
  # [8/8] unrolled `__torch_dispatch__` of `torch.spmm(t212, t205)`
  t219 = ltorch.core_aten_scaled_mm(t211, t216, t217, t218, None, None, torch.float32, True)  # t219: "cuda:0 f32[16, 64]"

(quote of https://gist.github.com/crcrpar/9b69ef83e68e306415af091d025cbf9c#file-1_1_fwd_trace_torchao_fp8tensor_flattened-py-L140-L158)
Here, neither torch._scaled_mm nor torch.ops.aten._scaled_mm have their backward defined in PyTorch. Neither thunder does.
So, this trace transform would better wait on the forward-backward split as at least we know the backward of torch.mm and ltorch.spmm.

Backward traces are also available in the linked gist.

I choose to apply the trace transform AFTER the forward-backward split after I worked on torchao.float8 linears.

Interesting!

  # /path/to/site-packages/torchao/float8/float8_linear.py:57: 	        res_bits = res_bits.reshape(*orig_shape[:-1], res_bits.shape[-1])
  tensor = manual_float8_matmul_with_args_in_float8_140062104971168_2(input_fp8, weight_fp8_t)  # tensor: "cuda:0 f32[16, 64]"
    # t91 = ltorch.reshape(input_fp8, -1, 32)  # t91: "cuda:0 f32[16, 32]"
      # t91 = prims.reshape(input_fp8, (16, 32))  # t91: "cuda:0 f32[16, 32]"
    # tensor = ltorch.spmm(t91, weight_fp8_t)  # tensor: "cuda:0 f32[16, 64]"

The corresponding lines are pytorch/ao@fe5f11b/torchao/float8/float8_linear.py#L32-L58, which is the forward of manual_float8_matmul_with_args_in_float8(torch.autograd.Function).

Here, t91 and weight_fp8_t are an object of torchao.float8.Float8Tensor

Following a previous comment I made, it would be great if they printed an annotation that showed they were tensor subclasses.

(quote of gist.github.com/crcrpar/9b69ef83e68e306415af091d025cbf9c#file-1_1_fwd_trace_torchao_fp8tensor_flattened-py-L140-L158) Here, neither torch._scaled_mm nor torch.ops.aten._scaled_mm have their backward defined in PyTorch. Neither thunder does. So, this trace transform would better wait on the forward-backward split as at least we know the backward of torch.mm and ltorch.spmm.

I was thinking it would be a good thing to get ops like torch.ops.aten._scaled_mm and then Thunder would have to define an autograd formula for them or a decomposition for them. Isn't that the point of adding the aten operators?

I already recreated a longer sequence of PRs: #1583, #1584, #1585, #1591, and #1592 and the last one improves the type_string.

Isn't that the point of adding the aten operators?

I don't think so. Some core aten ops don't define their backward nor decompositions like torch._scaled_mm and torch.ops.aten._scaled_mm.

I already recreated a longer sequence of PRs: #1583, #1584, #1585, #1591, and #1592 and the last one improves the type_string.

Isn't that the point of adding the aten operators?
I don't think so. Some core aten ops don't define their backward nor decompositions like torch._scaled_mm and torch.ops.aten._scaled_mm.

How does PyTorch generate the backward for those operations on the subclass, then?

The __torch_dispatch__ comes last as in https://pytorch.org/docs/stable/notes/extending.html#extending-torch-native-api:

For the purpose of extending torch, the important subset of the ordering for this discussion is:

vmap -> Autocast -> Autograd -> ZeroTensor -> Neg/Conj -> Functionalize -> Python -> Backends

autograd would not see torch._scaled_mm as it's used inside Float8Tensor.__torch_dispatch__.

The __torch_dispatch__ comes last as in pytorch.org/docs/stable/notes/extending.html#extending-torch-native-api:

For the purpose of extending torch, the important subset of the ordering for this discussion is:
vmap -> Autocast -> Autograd -> ZeroTensor -> Neg/Conj -> Functionalize -> Python -> Backends

autograd would not see torch._scaled_mm as it's used inside Float8Tensor.__torch_dispatch__.

Interesting!

OK; so to return to the sample the initial computation would look like:

# Trace
def computation(x):
  # x: "MyTensorSubclass[cuda:0 f32[12, 12]]" 
  t0 = MyTensorSubclass.torch.add(x, 1)  # t0: "MyTensorSubclass[cuda:0 f32[12, 12]]"
  return t0

And if that was differentiated then we'd ask the executors if they have a grad rule for MyTensorSubclass.torch.add and if not then we'd differentiate it like torch.add, except all the torch operations would be MyTensorSubclass variations?

Like if we had MyTensorSubclass.torch.mul, it would differentiate into MyTensorSubclass.torch.mul in the forward and MyTensorSubclass.torch.mul in the backward?

# Fwd + Bwd Trace
def computation(x):
  # x: "MyTensorSubclass[cuda:0 f32[12, 12]]" 
  t0 = MyTensorSubclass.torch.mul(x, 1)  # t0: "MyTensorSubclass[cuda:0 f32[12, 12]]"
  return t0
  # grad g for t0 introduced here
  t1 = MyTensorSubclass.torch.mul(g, x)
  return t1

And then we'd have a new transform, like "add torch dispatch decompositions" that would add the ATen operations under these?

So, assuming the subclass maps mul to div (which is silly, but whatever):

# After add torch dispatch decompositions (adds decompositions to the subclass)
def computation(x):
  # x: "MyTensorSubclass[cuda:0 f32[12, 12]]" 
  t0 = MyTensorSubclass.torch.mul(x, 1)  # t0: "MyTensorSubclass[cuda:0 f32[12, 12]]"
    # t0 = aten.div(x, 1)
  return t0
  # grad g for t0 introduced here
  t1 = MyTensorSubclass.torch.mul(g, x) # t1: "MyTensorSubclass[cuda:0 f32[12, 12]]"
    # t1 = aten.div(g, x)
  return t1

And then we could let executors claim these operations after that?

This ordering makes sense to me, I think. Is this what you had in mind, @crcrpar?

One thing that's unfortunate about this logic is that it doesn't let executors automatically apply their custom differentiation transformation to operations they could ultimately execute. For example, let's say someone writes a tensor subclass that maps to an aten operator that transformer engine could execute. If transformer engine had a custom fwd+bwd for that operation it can no longer apply it, because autodifferentiation has already happened. A fix for this is to register the actual tensor subclass operation with transformer engine.

Does that make sense? Does that fit your thinking, @kshitij12345, @IvanYashchuk?